from tomllib import load
import typing as T
import logging
from pydantic import BaseModel, Field, ConfigDict

from . import ROOT_DIR


class _DataSourceConfig(BaseModel):
    model_config = ConfigDict(extra="allow")
    dataname: str = Field(description="数据库类型")
    driver: str = Field(description="数据库驱动")
    username: str = Field(description="数据库用户名")
    password: str = Field(description="数据库密码")
    url: str = Field(description="数据库地址, host:prot/dbname 例如localhost:3306/test")
    echo: T.Optional[bool] = Field(
        default=False, description="是否开启sqlalchemy的日志"
    )
    pool_size: T.Optional[int] = Field(default=10, description="数据库连接池大小")
    pool_recycle: T.Optional[int] = Field(
        default=3600, description="数据库连接池回收时间, 单位秒"
    )
    max_overflow: T.Optional[int] = Field(
        default=25, description="允许超过池容量的最大连接数"
    )


class _RedisConfig(BaseModel):
    model_config = ConfigDict(extra="allow")
    host: str = Field(description="redis地址")
    port: int = Field(description="redis端口")
    password: T.Optional[str] = Field(default=None, description="redis密码")
    username: T.Optional[str] = Field(default=None, description="redis用户名")
    db: T.Optional[int] = Field(default=0, description="redis数据库")
    pool_size: T.Optional[int] = Field(default=10, description="redis连接池大小")
    time: T.Optional[int] = Field(default=43200, description="ttl时间, 单位秒")


class _CacheConfig(BaseModel):
    model_config = ConfigDict(extra="allow")
    update_time: T.Optional[int] = Field(
        default=3600, description="自动保存聊天记录时间, 单位秒"
    )
    max_cache: T.Optional[int] = Field(default=20, description="最多缓存的聊天记录数")
    max_use_history: T.Optional[int] = Field(
        default=10, description="最多使用的历史聊天记录数"
    )


class _MilvusConfig(BaseModel):
    model_config = ConfigDict(extra="allow")
    embedding_model_path: T.Optional[str] = Field(
        default="Alibaba-NLP/gte-multilingual-base",
        description="embedding模型路径或者名称",
    )
    host: str = Field(description="milvus地址")
    port: int = Field(description="milvus端口")
    db_name: T.Optional[str] = Field(default="default", description="milvus数据库名称")
    username: T.Optional[str] = Field(default=None, description="milvus用户名")
    password: T.Optional[str] = Field(default=None, description="milvus密码")
    pool_size: T.Optional[int] = Field(default=10, description="milvus连接池大小")


class ApplicationConfig(BaseModel):
    model_config = ConfigDict(extra="allow")
    datasource: _DataSourceConfig = Field(description="数据库配置")
    redis: _RedisConfig = Field(description="redis配置")
    cache: T.Optional[_CacheConfig] = Field(
        default=_CacheConfig(), description="缓存配置"
    )
    milvus: _MilvusConfig = Field(description="milvus配置")


def read():
    with open((ROOT_DIR / "config.toml").as_posix(), mode="rb") as file:
        return load(file)


configuration: ApplicationConfig = ApplicationConfig(**read())


logging.info("配置文件初始化完成...")

__all__ = ["configuration"]
